# -*- coding: utf-8 -*-
import urllib
import pybel

class PubChemObj:

  def __init__(self, cid=0):

    self.url = 'http://pubchem.ncbi.nlm.nih.gov/summary/summary.cgi'
    self.cid = cid
    self.dataXML = ''
    self.dataSDF = ''


  def getXML(self):

    location = urllib.urlretrieve(self.url+'?cid='+str(self.cid)+'&disopt=SaveXML')[0]

    dataFile = open(location, 'r')
    self.dataXML = dataFile.read()

  def writeXML(self, filename):

    if not self.dataXML:
      self.getXML()
    if not self.dataXML:
      return

    f = open(filename, 'w')
    f.write(self.dataXML)


  def getSDF(self):

    location = urllib.urlretrieve(self.url+'?cid='+str(self.cid)+'&disopt=SaveSDF')[0]

    dataFile = open(location, 'r')
    self.dataSDF = dataFile.read()

  def writeSDF(self, filename):

    if not self.dataXML:
      self.getXML()
    if not self.dataXML:
      return

    f = open(filename, 'w')
    f.write(self.dataXML)

  def getMol(self):

    return pybel.readstring('sdf', self.dataSDF)
